// Last Update:2015-04-15 09:38:22
/**
 * @file main.c
 * @brief
 * @author
 * @version 0.1.00
 * @date 2015-01-26
 */

#include <errno.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include <unistd.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>

#include <arpa/inet.h>
#include <net/if.h>
#include <net/ethernet.h>
#include <net/ethernet.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/in.h>

#include <pcap.h>
#include <zlog.h>
#include <libHX/io.h>
#include <libHX/option.h>

#define LOG_CONF "/dev/shm/.logformat.conf"
#define DEV_LEN 8
#define MAC_LEN 6
#define PATH_LEN 256
#define IP_LEN 16
#define CMD_LEN 256
#define MAX_BUFLEN 128
#define VERSION "0.1"

struct options
{
	char path[PATH_LEN];
	char dev[DEV_LEN];
	uint8_t smac[MAC_LEN];
	uint8_t dmac[MAC_LEN];
	char sip[IP_LEN];
	char dip[IP_LEN];
	uint16_t sport;
	uint16_t dport;
};
struct options g_options;
zlog_category_t *zc = NULL;
int is_save;

int log_init();
int get_eth_mac(char *if_name, uint8_t *smac);
void get_eth_ip(char *if_name, char *sip);
int option_init(int argc, char *argv[]);
void opt_version(const struct HXoptcb *cbi);
void parse_mac(char *if_name, char *type, uint8_t *mac, uint8_t *dst_mac);
void get_gateway_mac(char *if_name, uint8_t *mac);
char *get_eths();
int get_packet(pcap_t *pcap_descr, const u_char **packet, struct pcap_pkthdr *hdr);
uint16_t tcp_checksum(struct iphdr *iph, int iphdr_len,
		struct tcphdr *tcph, int tcphdr_len, int payload_len);
void process_packet(const u_char *packet, struct pcap_pkthdr hdr,
		uint8_t *smac, char *sip, uint8_t *dmac, char *dip);
uint16_t ip_checksum(struct iphdr *iph, int iphdr_len);


int log_init()
{
	char *log_rules = "[rules]\n*.* \"/dev/shm/capwrite.log\",10M*5;\n";
	FILE *fp = fopen(LOG_CONF, "wr");
	if (fp == NULL)
	{
		perror("Log init error");
		return -1;
	}
	fputs(log_rules, fp);
	fclose(fp);

	int rc;
	rc = zlog_init(LOG_CONF);
	if (rc)
	{
		fprintf(stderr, "Zlog init failed\n");
		return -1;
	}
	zc = zlog_get_category("my_cat");
	if (!zc)
	{
		fprintf(stderr, "Get log handle error\n");
		zlog_fini();
		return -1;
	}

	return 0;
}

int get_packet(pcap_t *pcap_descr, const u_char **packet, struct pcap_pkthdr *hdr)
{
	*packet = pcap_next(pcap_descr, hdr);
	if (*packet == NULL)
	{
		return 0;
	}

	return 1;
}

uint16_t tcp_checksum(struct iphdr *iph, int iphdr_len,
		struct tcphdr *tcph, int tcphdr_len, int payload_len)
{
	uint16_t len, i;
	uint32_t sum = 0;

	uint8_t *data;
	data = (uint8_t *)(&iph->saddr) ;
	for(i = 0; i < 8; i += 2)
	{
		sum += *(uint16_t *)data;
		data +=2;
	}

	len = tcphdr_len + payload_len;
	sum += htons(len + iph->protocol);
	data = (uint8_t *)tcph;

	if(len & 0x1) {
		data[len] = 0;
		len ++;
	}
	for(i = 0; i < len; i += 2)
	{
		sum += *(uint16_t *)data;
		data += 2;
	}
	sum = (sum >> 16) + (sum & 0xffff);
	sum += (sum >> 16);
	return (uint16_t)(~sum & 0xFFFF);
}


uint16_t ip_checksum(struct iphdr *iph, int iphdr_len)
{
	uint16_t len, i;
	uint32_t sum = 0;

	uint8_t *data;
	len = iphdr_len;
	data = (uint8_t *)(iph) ;
	if (len & 0x1)
	{
		data[len] = 0;
		len ++;
	}
	for(i = 0; i < len; i += 2)
	{
		sum += *(uint16_t *)data;
		data += 2;
	}
	sum = (sum >> 16) + (sum & 0xffff);
	sum += (sum >> 16);

	return (uint16_t)(~sum & 0xFFFF);
}

void process_packet(const u_char *packet, struct pcap_pkthdr hdr,
		uint8_t *smac, char *sip, uint8_t *dmac, char *dip)
{
	struct ether_header *eptr = (struct ether_header *)packet;
	struct iphdr *iph = NULL;
	struct tcphdr *tcph = NULL;
	int eth_length = 14;
	int iph_length = 20;
	int tcph_length = 20;
	int payload_len = hdr.caplen - (eth_length + iph_length + tcph_length);
	zlog_debug(zc, "-----------------------------------------------------");
	zlog_debug(zc, "Packet length: %d", hdr.caplen);
	zlog_debug(zc, "Payload length: %d", payload_len);

	memcpy(eptr->ether_shost, smac, 6);
	memcpy(eptr->ether_dhost, dmac, 6);

	switch (ntohs(eptr->ether_type))
	{
		case ETHERTYPE_VLAN:
			{
				uint16_t ether_type = *((uint16_t*)(eptr + 1) + 1);
				if (ntohs(ether_type) != ETHERTYPE_IP)
				{
					return ;
				}
				iph = (struct iphdr*)((uint16_t*)(eptr + 1) + 2);
				if (iph->protocol == IPPROTO_TCP)
				{
					tcph = (struct tcphdr *)(iph + iph_length);
					if (g_options.sport != 0)
						tcph->source= htons(g_options.sport);
					if (g_options.dport != 0)
						tcph->dest = htons(g_options.dport);

					goto cal;
				}
			}
		case ETHERTYPE_IP:
			{
				iph = (struct iphdr*)(eptr + 1);
				if (iph->protocol ==  IPPROTO_TCP)
				{
					tcph = (struct tcphdr *)((uint8_t *)iph + iph_length);
					if (g_options.sport != 0)
						tcph->source = htons(g_options.sport);
					if (g_options.dport != 0)
						tcph->dest = htons(g_options.dport);
					goto cal;
				}
				return ;
			}
		case ETHERTYPE_ARP:
			return ;
		default:
			return ;
	}

cal:
	iph->saddr = inet_addr(sip);
	if (dip != NULL)
		iph->daddr = inet_addr(dip);
	iph->check = 0;
	iph->check = ip_checksum(iph, iph_length);
	zlog_debug(zc, "ip->checksum:  %02x   ", iph->check);
	tcph->check = 0;
	tcph->check = tcp_checksum(iph, iph_length, tcph, tcph_length, payload_len);
	zlog_debug(zc, "tcp->checksum: %02x\n", tcph->check);
}

void get_pcap_files(char *dirname, char **files)
{
	const char *dentry = NULL;
	struct HXdir *dh;
	if ((dh = HXdir_open(dirname)) == NULL)
	{
		fprintf(stderr, "Could not open directory: %s\n", strerror(errno));
		return;
	}
	int i = 0;
	char path[80] = {0};
	realpath(dirname, path);
	zlog_debug(zc, "realpath of %s: %s", dirname, path);
	while ((dentry = HXdir_read(dh)) != NULL)
	{
		struct stat sb;
		memset(&sb, 0, sizeof(sb));
		files[i] = malloc(256);
		memset(files[i], 0, 256);
		sprintf(files[i], "%s/%s", path, dentry);
		if (stat(files[i], &sb) == -1)
		{
			perror("Get files stat error");
			exit(EXIT_SUCCESS);
		}
		switch (sb.st_mode & S_IFMT)
		{
			case S_IFREG:
				break;
			default:
				continue;
		}
		i++;
	}
	files[i] = NULL;
	HXdir_close(dh);
}


char *get_eths()
{
	char *dev = NULL;

	char errbuf[PCAP_ERRBUF_SIZE] = {0};
	dev = pcap_lookupdev(errbuf);
	if (dev == NULL)
	{
		fprintf(stderr, "Lookupdev error: %s\n", errbuf);
		exit(EXIT_FAILURE);
	}
	return dev;

}


void get_gateway_mac(char *if_name, uint8_t *mac)
{
	char cmd[CMD_LEN] = {0};
	char output[MAX_BUFLEN] = {0};
	snprintf(cmd, CMD_LEN - 1, "arp -a `route -n | /bin/grep UG | /bin/grep \'%s\' |\
			awk \'{print $2}\'` | cut -d \" \" -f 4", if_name);
	zlog_debug(zc, "get_eth_ip, %s: ", cmd);
	FILE *fp = popen(cmd, "r");
	if (fp == NULL)
	{
		fprintf(stderr, "Get gateway ip error, please set dip \n");
	}
	if (fgets(output, MAX_BUFLEN, fp) != NULL)
	{
		sscanf((const char *)output, "%02x%*[:-]%02x%*[:-]%02x%*[:-]%02x%*[:-]%02x%*[:-]%02x",
				(unsigned int *)&mac[0], (unsigned int *)&mac[1], (unsigned int *)&mac[2],
				(unsigned int *)&mac[3], (unsigned int *)&mac[4], (unsigned int *)&mac[5]);
	}
	pclose(fp);
}



void parse_mac(char *if_name, char *type, uint8_t *mac, uint8_t *dst_mac)
{
	if (mac == NULL)
	{
		if (strcmp(type, "LOCAL") == 0)
			get_eth_mac(if_name, dst_mac);
		else if (strcmp(type, "GATEWAY") == 0)
			get_gateway_mac(if_name, dst_mac);
	}
	else
	{
		if (strlen((const char *)mac) != 17)
		{
			fprintf(stderr, "Mac is not correct\n");
			exit(EXIT_FAILURE);
		}

		sscanf((const char *)mac, "%02x%*[:-]%02x%*[:-]%02x%*[:-]%02x%*[:-]%02x%*[:-]%02x",
				(unsigned int *)&dst_mac[0], (unsigned int *)&dst_mac[1], (unsigned int *)&dst_mac[2],
				(unsigned int *)&dst_mac[3], (unsigned int *)&dst_mac[4], (unsigned int *)&dst_mac[5]);
	}
}

void opt_version(const struct HXoptcb *cbi)
{
	fprintf(stdout, "Capwrite-%s, 28 Jan 2015\n", (const char *)cbi->current->uptr);
    exit(EXIT_SUCCESS);
}

int option_init(int argc, char *argv[])
{
	char *dev = NULL;
	char *smac = NULL;
	char *dmac = NULL;
	char *sip = NULL;
	char *dip = NULL;
	char *path = NULL;
	char *sport = NULL;
	char *dport = NULL;
	struct  options opts;
	memset(&opts, 0, sizeof(opts));
	memset(&g_options, 0, sizeof(g_options));

    struct HXoption options_table[] = {
        {.sh = 'i', .ln = "interface", .type = HXTYPE_STRING, .ptr = &dev,     .help = "Set interface of sending packets"},
        {.sh = 'd', .ln = "dir",       .type = HXTYPE_STRING, .ptr = &path,    .help = "Directory of saving pcaps"},
        {.sh = 's', .ln = "save",      .type = HXTYPE_NONE,   .ptr = &is_save, .help = "A switch for saving pcaps"},
        {           .ln = "smac",      .type = HXTYPE_STRING, .ptr = &smac,    .help = "Set Source MAC"},
        {           .ln = "dmac",      .type = HXTYPE_STRING, .ptr = &dmac,    .help = "Set Destination MAC"},
        {           .ln = "sip",       .type = HXTYPE_STRING, .ptr = &sip,     .help = "Set Source IP"},
        {           .ln = "dip",       .type = HXTYPE_STRING, .ptr = &dip,     .help = "Set Destination IP"},
		{			.ln = "sport",	   .type = HXTYPE_STRING, .ptr = &sport,   .help = "Set Source Port"},
		{			.ln = "dport",	   .type = HXTYPE_STRING, .ptr = &dport,   .help = "Set Destination Port"},
        {.sh = 'v', .ln = "version",   .type = HXTYPE_NONE,   .uptr = VERSION, .help = "Show version information", .cb = opt_version},
        HXOPT_AUTOHELP,
        HXOPT_TABLEEND,
    };
	if (HX_getopt(options_table, &argc, (const char ***)&argv, HXOPT_USAGEONERR) !=
			HXOPT_ERR_SUCCESS)
		return EXIT_FAILURE;

	if (dev == NULL)
		strcpy(g_options.dev, get_eths());
	else
		strcpy(g_options.dev, dev);
		
	parse_mac(g_options.dev, "LOCAL",   (uint8_t *)smac, g_options.smac);
	parse_mac(g_options.dev, "GATEWAY", (uint8_t *)dmac, g_options.dmac);
	if (path == NULL)
	{
		printf("Must specified a pcaps directory, use -d [directory].\n");
		return -1;
	}
	else
		strcpy(g_options.path, path);

	if (sip  == NULL)
		get_eth_ip(g_options.dev, g_options.sip);
	else
		strcpy(g_options.sip, sip);
	if (dip != NULL)
		strcpy(g_options.dip, dip);
	(sport == NULL) ? (g_options.sport = 0) : (g_options.sport = atoi(sport));
	(dport == NULL) ? (g_options.dport = 0) : (g_options.dport = atoi(dport));
	zlog_debug(zc, "option: %s, %s, %s, %s, %s, %s\n", dev, path, smac, dmac, sip, dip);

	return EXIT_SUCCESS;
}

void get_eth_ip(char *if_name, char *sip)
{
	struct sockaddr_in sin;
	struct ifreq ifr;
	int sockfd = socket(AF_INET, SOCK_DGRAM, 0);

	memset(&sin, 0, sizeof(sin));
	memset(&ifr, 0, sizeof(ifr));

	if (if_name == NULL)
	{
		fprintf(stderr, "%s\n", "Not specified the NIC.");
		exit(EXIT_FAILURE);
	}
	/* The device that we want to get information */
	strncpy(ifr.ifr_name, if_name, IFHWADDRLEN);

	/* Get IP Address */
	if ( (ioctl(sockfd, SIOCGIFADDR, &ifr)) < 0)
	{
		perror("Can not get IP Address");
	}
	memcpy(&sin, &ifr.ifr_addr, sizeof(sin));
	strcpy(sip, inet_ntoa(sin.sin_addr));
	printf("sip: %s\n", sip);
}

int get_eth_mac(char *if_name, uint8_t *smac)
{
	struct sockaddr_in sin;
	struct ifreq ifr;
	int sockfd = socket(AF_INET, SOCK_DGRAM, 0);

	memset(&sin, 0, sizeof(sin));
	memset(&ifr, 0, sizeof(ifr));

	if (if_name == NULL)
	{
		fprintf(stderr, "missing arguments\n");
		exit(EXIT_FAILURE);
	}
	/* The device that we want to get information */
	strncpy(ifr.ifr_name, if_name, IFHWADDRLEN);

	/* Get the Hardware Address */
	if ( (ioctl(sockfd, SIOCGIFHWADDR, &ifr)) < 0)
	{
		perror("Can not get Hardware Address");
	}

	memcpy(smac, ifr.ifr_hwaddr.sa_data, 6);
	return 0;
}

int main(int argc, char *argv[])
{
	char *files[256];
	char error_buf[PCAP_ERRBUF_SIZE] = {0};
	memset(files, 0, sizeof(files));

	if (option_init(argc, argv) != 0)
	{
		exit(EXIT_FAILURE);
	}
	if (log_init() != 0)
	{
		exit(EXIT_FAILURE);
	}
	printf("dev:%s, path:%s, sip:%s, dip:%s, ",
			g_options.dev, g_options.path, g_options.sip, g_options.dip);
	printf("smac: ");
	int i;
	for (i = 0; i < 6; i++)
	{
		printf("%02x:", g_options.smac[i]);
	}
	printf(", dmac: ");
	for (i = 0; i < 6; i++)
	{
		printf("%02x:", g_options.dmac[i]);
	}
	printf("\n");
	get_pcap_files(g_options.path, files);

	pcap_t *pv = pcap_open_live(g_options.dev, 100, 0, -1, error_buf);
	if (pv == NULL)
	{
		fprintf(stderr, "pcap open live error: %s\n", error_buf);
		zlog_debug(zc, "pcap_open_live error: %s\n", error_buf);
		exit(EXIT_FAILURE);
	}

	for (i = 0; files[i] != NULL ; i++)
	{
		printf("Processing %s\n", files[i]);
		zlog_debug(zc, "Processing %s", files[i]);
		pcap_t *pcap_desc = pcap_open_offline(files[i], error_buf);
		const u_char *packet = NULL;
		struct pcap_pkthdr hdr; /*  pcap.h */

		char outfile[256] = {0};
		sprintf(outfile, "/dev/shm/%d-cache.cap", i);
		zlog_debug(zc, "Cache file: %s", outfile);
		pcap_dumper_t *pout = pcap_dump_open(pcap_desc, outfile);
		while (1)
		{
			if (get_packet(pcap_desc, &packet, &hdr) > 0)
			{
				process_packet(packet, hdr, g_options.smac, g_options.sip, g_options.dmac, g_options.dip);
				if (is_save)
				{
					pcap_dump((u_char *)pout, &hdr, packet);
				}
				pcap_sendpacket(pv, packet, hdr.len);
			}
			else
			{
				break;
			}
		}
		printf("Send packet success\n");
		zlog_debug(zc, "Send packet success");
		pcap_dump_close(pout);
		pcap_close(pcap_desc);
		free(files[i]);
	}
	zlog_fini();

	return 0;
}
